#!/usr/bin/env python
#-*- coding: utf-8 -*-

import numpy as np
import matplotlib.pyplot as plt

import pat.dynamic_model as dm
import pat.autopilot     as ap
import pat.vehicles.rotorcraft.autopilot_planar as app
import pat.utils as pu

my_dm = dm.load(type="rotorcraft", variant="planar_simple")
my_ap = ap.load(type="rotorcraft", variant="planar", _dm=my_dm)

print "Name: ", my_dm.name()
print "Param:\n", my_dm.param()
Xe, Ue = my_dm.trim()
print "Xe: ", Xe
print "Ue: ", Ue

my_ap.set_cur_mode(7) # PositionLawOpt
my_ap.set_ref(app.Step_x(Xe, Ue))

X0 = np.array(Xe)
X0[my_dm.sv_x]+= 1.
X0[my_dm.sv_z]+= 1.
X0[my_dm.sv_th]+= pu.rad_of_deg(1.)
my_dm.reset(X0)
dt = 0.01; time = np.arange(0.,20., dt)
X = np.zeros((len(time), my_dm.sv_size)); X[0] = X0
U = np.zeros((len(time), my_dm.iv_size));
for i in range(1,len(time)):
    U[i-1], Refi = my_ap.run(time[i-1], X[i-1])
    X[i] = my_dm.run(dt, U[i-1])
U[-1], refi = my_ap.run(time[-1], X[-1])

my_dm.plot_state(time, X, U, window_title="LQR position law", filename="../doc/images/wiki/planar_vtol_state_feedback_4.png")
plt.show()
